import argparse
import http.server
import json
import os
import socketserver

import mss
import pyautogui
from PIL import Image
from ultralytics import YOLO

wd = os.getcwd()


class ScreenCapture:

    def __init__(self):
        w, h = pyautogui.size()
        self.sct = mss.mss()  # 实例化mss，并使用高效模式 mon=-1, optimize=True
        self.monitor_settings = {
            'left': 0,
            'top': 0,
            'width': w,
            'height': h
        }

    def capture(self):
        img = self.sct.grab(self.monitor_settings)
        return Image.frombytes('RGB', img.size, img.bgra, 'raw', 'BGRX')


def getModel():
    model = YOLO(os.path.join(wd, "model.yaml"))
    model = YOLO(os.path.join(wd, "best.pt"))
    return model


def train(model: YOLO):
    model.train(data=os.path.join(wd, "data.yaml"), epochs=4000, imgsz=1280, patience=100, batch=-1, device=0,
                workers=3)
    # metrics = model.val()
    # success = model.export(format="onnx", opset=12, simplify=True, dynamic=False, imgsz=640)
    return model


defaultModel = getModel()
sc = ScreenCapture()


def capture_predict(model: YOLO, scap: ScreenCapture):
    try:
        img = scap.capture()
        results = model(source=img, imgsz=1280, save=False, name=os.path.join(wd, "runs/detect"))

        # _, binary_image = cv2.threshold(cv2.cvtColor(np.array(img), cv2.COLOR_RGB2GRAY), 1, 255, cv2.THRESH_BINARY)
        # black_pixels = np.sum(binary_image == 0)
        # black_ratio = round(black_pixels/binary_image.size, 2)
        # if black_ratio > 0.4:
        #     return {
        #         "via_map": [{"label": "via_map", "confidence": 0.9}]
        #     }

        r = {}
        if len(results) == 0:
            return r

        boxes = results[len(results) - 1:][0]
        names = boxes.names
        boxes = boxes.boxes

        for i, cls in enumerate(boxes.cls):
            label = names[int(cls)]
            confidence = round(float(boxes.conf[i]), 2)
            x = int(boxes.xywh[i][0])
            y = int(boxes.xywh[i][1])

            if ['character', 'open_door', 'close_door', 'money', 'item', 'epic', 'guai_wu'].__contains__(label):
                y = int(boxes.xyxy[i][3])

            dr = {"label": label, "confidence": confidence, "x": x, "y": y}
            if r.__contains__(label):
                r[label].append(dr)
            else:
                r.setdefault(label, [dr])

        return r
    except Exception as e:
        print(f"{e}")


class HttpHandler(http.server.SimpleHTTPRequestHandler):
    def do_GET(self):
        try:
            idx = self.path.index("?")
            if idx != -1:
                msg = self.path[idx + 1:].split("=")[1]
                print(msg)

                if msg == "":
                    self.send_response(500)
                    return
                else:
                    if msg == "ping":
                        self.send_response(200)
                        self.send_header('Content-type', 'application/json')
                        self.end_headers()
                        self.wfile.write(b"")

                    if msg == "capture":
                        r = capture_predict(defaultModel, sc)

                        self.send_response(200)
                        self.send_header('Content-type', 'application/json')
                        self.end_headers()
                        self.wfile.write(json.dumps(r).encode("utf-8"))

        except:
            self.send_response(200)
            self.send_header('Content-type', 'application/json')
            self.end_headers()
            self.wfile.write(json.dumps({}).encode("utf-8"))


def start_server():
    parser = argparse.ArgumentParser()
    parser.add_argument('--port', type=int, default=8848)
    opt = parser.parse_args()

    with socketserver.TCPServer(('127.0.0.1', opt.port), HttpHandler) as httpd:
        print('Server listening on port', opt.port)
        capture_predict(defaultModel, sc)
        httpd.serve_forever()


if __name__ == '__main__':
    train(defaultModel)
    # start_server()

    # results = defaultModel(source=r'E:\ideaworkspace3\ultralytics\VOCdevkit\images\train', imgsz=1280,
    #                        save=True, name=os.path.join(wd, "runs/detect"))


